from typing import List, Dict, Any
import threading
from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn
from transformers import AutoTokenizer
from src.index_tree import IndexTree
from src.services.zmq_manager import ZMQManager
from src.logging import get_logger
from src.config import get_config
from src.services.service_discovery import K8sServiceDiscovery


class MatchSortRequest(BaseModel):
    """
    用于接收/match_sort接口的请求体。
    server_ip: server_ip:port字符串列表
    body: 包含model和messages的字典
    """
    server_ip: List[str]  # ["127.0.0.1:8000", "127.0.0.1:8001", ...]
    body: Dict[str, Any]  # {"model": "model_name", "messages": [...]}


class MatchSortResponse(BaseModel):
    """
    /match_sort接口的响应体
    server_score_list: 匹配结果列表
    message: 提示信息
    status: 业务状态码
    """
    server_score_list: List[Dict[str, float]]
    message: str = ""
    status: int = 0


class TokenizerManager:
    """
    Tokenizer管理器，负责tokenizer的加载和缓存
    """
    def __init__(self):
        self.tokenizer_cache: Dict[str, AutoTokenizer] = {}
        self.logger = get_logger("tokenizer_manager")

    def _get_tokenizer(self, model_name: str) -> AutoTokenizer:
        """
        获取tokenizer，如果缓存中没有则加载并缓存
        :param model_name: 模型名称
        :return: AutoTokenizer实例
        """
        if model_name not in self.tokenizer_cache:
            self.logger.info(f"首次加载tokenizer:{model_name}")
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            self.tokenizer_cache[model_name] = tokenizer
        else:
            tokenizer = self.tokenizer_cache[model_name]
            self.logger.info(f"使用缓存的tokenizer:{model_name}")
        return tokenizer

    def get_tokens(self, model_name: str, messages: List[Dict]) -> List[int]:
        """
        将messages转换为tokens
        :param model_name: 模型名称
        :param messages: 消息列表
        :return: token列表
        """
        tokenizer = self._get_tokenizer(model_name)
        tokens = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True)
        self.logger.info(f"token转换完毕")
        self.logger.debug(f"tokens: {tokens}")
        return tokens
        

class KVIndexer:
    """
    KVIndexer 负责管理几群内 vllm server 的 kv_cache 缓存信息的前缀树。
    属性：
        server_list: 通过服务发现获取的服务器列表
        ip_forest_dict: 每个 server_ip:port 对应的 IndexTree 列表
        tokenizer_manager: tokenizer管理器
    """
    def __init__(self):
        """
        初始化KVIndexer，完成服务发现、数据结构初始化，并启动对router的http接口服务。
        """
        self.logger = get_logger("kv_indexer")
        self.logger.info("Initializing KVIndexer")
        
        # 初始化tokenizer管理器
        self.tokenizer_manager = TokenizerManager()
        
        # 初始化服务发现相关属性
        self.server_ip_port_map: Dict[str, str] = {}
        self.ip_forest_dict: Dict[str, List[IndexTree]] = {}
        self.server_list: List[Dict[str, str]] = []
        self.discovery_thread = None
        self.discovery_running = False

        # 启动服务发现线程
        self._start_discovery_thread()
        
        # 第二部分：初始化数据结构        
        self.logger.info("Server discovery refresh, current servers: ",
                        endpoint_count=len(self.ip_forest_dict),
                        endpoints=list(self.ip_forest_dict.keys()))
        
        self.zmq_manager = ZMQManager(self.ip_forest_dict)
        self.zmq_manager.start()
        
        # 第三部分：初始化KVIndexerService
        self.service = KVIndexerService(self)

    def _start_discovery_thread(self):
        """
        启动服务发现线程，定期刷新服务列表
        """
        self.discovery_running = True
        self.discovery_thread = threading.Thread(target=self._discovery_loop, daemon=True)
        self.discovery_thread.start()
        self.logger.info("Service discovery thread started")

    def _discovery_loop(self):
        """
        服务发现循环，根据配置的间隔时间执行
        """
        import time
        
        config = get_config()  # 使用顶层导入的get_config
        refresh_interval = config.service_discovery.refresh_interval
        
        while self.discovery_running:
            try:
                # 执行服务发现
                curr_server_list: List[Dict[str, str]] = self._discover_servers()

                new_server_ip_map = {}
                for server_dict in curr_server_list:
                    for ip, port in server_dict.items():
                        key = f"{ip}:{port}"
                        if key not in self.ip_forest_dict:
                            new_server_ip_map[ip] = port

                remove_server_ip_map = {}
                for server_dict in self.server_list:
                    for ip, port in server_dict.items():
                        if {ip: port} not in curr_server_list:
                            remove_server_ip_map[ip] = port

                self._update_server_info(new_server_ip_map, remove_server_ip_map)
                self.server_list = curr_server_list

                # 根据配置的间隔时间等待
                time.sleep(refresh_interval)
                
            except Exception as e:
                self.logger.error("Service discovery loop error", error=str(e), exc_info=True)
                time.sleep(refresh_interval) 

    def _update_server_info(self, new_server_ip_map: Dict[str, str], remove_server_ip_map: Dict[str, str]):
        """
        更新ip_forest_dict, server_ip_port_map，并同步更新ZMQ worker
        """
        for ip, port in new_server_ip_map.items():
            key = f"{ip}:{port}"
            self.ip_forest_dict[key] = []
            self.server_ip_port_map[ip] = port
            # 添加新的ZMQ worker
            self.zmq_manager.add_worker(key)
            self.logger.info("Added new server and ZMQ worker", ip=ip, port=port, key=key)

        for ip, port in remove_server_ip_map.items():
            key = f"{ip}:{port}"
            # 先移除ZMQ worker
            self.zmq_manager.remove_worker(key)
            # 再移除数据结构
            self.ip_forest_dict.pop(key, None)
            self.server_ip_port_map.pop(ip, None)
            self.logger.info("Removed server and ZMQ worker", ip=ip, port=port, key=key)

    def stop_discovery(self):
        """
        停止服务发现线程
        """
        self.discovery_running = False
        if self.discovery_thread and self.discovery_thread.is_alive():
            self.discovery_thread.join(timeout=5)
        self.logger.info("Service discovery thread stopped")

    def _discover_servers(self) -> List[Dict[str, str]]:
        """
        服务发现功能，返回服务器列表。
        :return: 服务器列表，每个元素为{ip: port}的字典
        """
        try:
            # 读取配置
            config = get_config()
            service_discovery_config = config.service_discovery
            namespace = service_discovery_config.namespace
            label_selector = service_discovery_config.label_selector
            port_name = service_discovery_config.port_name
            
            self.logger.info("Starting K8s pod discovery",
                           namespace=namespace,
                           label_selector=label_selector,
                           port_name=port_name)
            
            # 初始化K8s服务发现
            k8s_discovery = K8sServiceDiscovery(
                namespace=namespace,
                label_selector=label_selector
            )
            
            # 获取服务列表
            services = k8s_discovery.refresh_services()
            
            # 转换为所需格式
            servers = []
            for service in services:
                    
                ip = service.get('ip')
                port = service.get('port')
                
                if ip and port:
                    servers.append({ip: port})
                    self.logger.info("Found pod", ip=ip, port=port)
            
            self.logger.info("K8s service discovery completed", server_count=len(servers))
            return servers
            
        except Exception as e:
            self.logger.error("K8s service discovery failed", error=str(e), exc_info=True)
            return []

    def match_sort(self, server_ip_list: List[str], tokens: List[int]) -> List[Dict[str, float]]:
        """
        对指定的server_ip列表，计算每个节点的匹配率，并按匹配率降序排序。
        :param server_ip_list: 需要匹配的server_ip:port字符串列表
        :param tokens: 查询用的token列表
        :return: [{server_ip: match_ratio}, ...] 按match_ratio降序排列
        """
        self.logger.info("Starting match_sort",
                        server_count=len(server_ip_list),
                        token_count=len(tokens))
        
        result = []
        for ip in server_ip_list:
            match_ratio = 0.0
            if ip not in self.server_ip_port_map:
                self.logger.warn("Server IP not found in server_ip_port_map", server=ip)
                continue
            port = self.server_ip_port_map.get(ip)
            ip_port = f"{ip}:{port}"
            forest: List[IndexTree] = self.ip_forest_dict.get(ip_port, [])
            
            if not forest:
                self.logger.warn("No trees found for server", server=ip_port)
            
            for tree in forest:
                match_count = tree.match_token_size(tokens)
                if len(tokens) > 0:
                    match_ratio = round(match_count / len(tokens), 2)
                if (match_ratio != 0):
                    self.logger.debug("Found match",
                                   server=ip_port,
                                   match_ratio=match_ratio,
                                   match_count=match_count)
                    break
            result.append({ip: match_ratio})
        
        # 按照match_ratio降序排序
        result.sort(key=lambda x: list(x.values())[0], reverse=True)
        
        self.logger.info("Match sort completed",
                        result_count=len(result),
                        best_match=result[0] if result else None)
        return result


class KVIndexerService:
    """
    KVIndexerService 封装FastAPI服务，负责注册路由并启动服务。
    """
    def __init__(self, kv_indexer: KVIndexer):
        self.kv_indexer = kv_indexer
        self.app = FastAPI()
        self.logger = get_logger("kv_indexer_service")
        self._register_routes()
        self.server_thread = None
        self.logger.info("KVIndexerService initialized")

    def _register_routes(self):
        @self.app.post("/match_sort", response_model=MatchSortResponse)
        async def match_sort_api(req: MatchSortRequest):
            self.logger.debug("Received match_sort request",
                           server_count=len(req.server_ip),
                           body=req.body)
            
            model_name = req.body.get("model")
            messages = req.body.get("messages", [])
            tokens = self.kv_indexer.tokenizer_manager.get_tokens(model_name, messages)
            
            result = self.kv_indexer.match_sort(req.server_ip, tokens)
            
            self.logger.info("Returning match_sort response",
                           result_count=len(result))
            return MatchSortResponse(server_score_list=result, message="", status=0)

    def start_service(self, host="0.0.0.0", port=28080, blocking=False):
        """启动HTTP服务
        
        Args:
            host: 服务绑定的主机地址
            port: 服务绑定的端口
            blocking: 如果为True，在主线程中阻塞运行；如果为False，在守护线程中运行
        """
        if blocking:
            self.logger.info("Starting HTTP service in main thread (blocking)", host=host, port=port)
            uvicorn.run(self.app, host=host, port=port)
        else:
            def run_server():
                self.logger.info("Starting HTTP service in thread", host=host, port=port)
                uvicorn.run(self.app, host=host, port=port)
            
            self.server_thread = threading.Thread(target=run_server, daemon=True)
            self.server_thread.start()
            self.logger.info("HTTP service thread started", host=host, port=port)
